#include "decode.h"
#include <bitstream.h>
#include "deembed.h"
#include <hwt.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <huffman/huffman.h>
//#include <libjbig/jbig.h>
#include <jbig.h>
#include <md5/md5.h>

int tian_grayscale_decode(
			  unsigned char **restored_data,
			  unsigned char **data,
			  unsigned int height,
			  unsigned int width,
			  const int verbose
			  )
{

#ifdef DEBUG
  unsigned int checksum;
#endif

  int i, j;

  /*
    1. compute difference-bitmap and average-bitmap (integer Haar wavelet transform)
  */
  if(verbose) printf("> compute integer Haar wavelet transform\n");
  
  //allocate the difference-bitmap
  int **differences = (int **)malloc(height * sizeof(int *));
  for(i=0; i<height; ++i)
    differences[i] = (int *)malloc(width/2 * sizeof(int));

  //allocate the average-bitmap
  unsigned char **average = (unsigned char **)malloc(height * sizeof(unsigned char *));
  for(i=0; i<height; ++i)
    average[i] = (unsigned char *)malloc(width/2 * sizeof(unsigned char));

  hwt_bitmap(differences, average, data, height, width);
#ifdef DEBUG
  int checksum_ = 0;
  checksum = 0;
  for(i=0; i<height; ++i)
    for(j=0; j<width/2; ++j) {
      checksum_ += differences[i][j];
      checksum += average[i][j];
    }
  printf("+ (emb.) differences checksum: %d\n", checksum_);
  printf("+ (emb.) averages checksum: %d\n", checksum);
#endif


  /*
    2. compute CH/NC membership-bitmap
  */
  if(verbose) printf("> classification\n");

  //allocate the membership-bitmap
  difference_class **membership = (difference_class **)malloc(height * sizeof(difference_class *));
  for(i=0; i<height; ++i)
    membership[i] = (difference_class *)malloc(width/2 * sizeof(difference_class));

  class_count cc;

  diff_classification_CH(membership, &cc, differences, average, height, width/2);

#ifdef DEBUG
  printf("|CH| = %d bits\n", cc.n_CH);
#endif


  /*
    3. extract the bitstream
  */
  if(verbose) printf("> extract the bitstream\n");
  unsigned char *bitstream_ = (unsigned char *)malloc(cc.n_CH * sizeof(unsigned char));
  unsigned char *bitstream = bitstream_;
  extract_bitstream(bitstream, differences, membership, width, height/2);

#ifdef DEBUG
  checksum = 0;
  for(i=0; i<cc.n_CH; ++i)
    checksum += bitstream[i];
  printf("bitstream checksum: %d\n", checksum);
#endif


  /*
    4.1. decode the location map
  */
  if(verbose) printf("> decode the (JBIG-compressed) location-map\n");
  unsigned int clmbs_length_wh = to_int32(bitstream, 0) * 8; //bits (inc. 32-bit header)

#ifdef DEBUG
  printf("|L| = %u bits\n", clmbs_length_wh);
#endif

  //decode JBIG
  if(verbose) printf("> [JBIG] uncompress the location-map\n");
  unsigned int clmbs_size_wh; //bytes (including 32-bit header)
  unsigned char *clmbs = to_bitstream_compact(bitstream, clmbs_length_wh, &clmbs_size_wh);
#ifdef DEBUG
  printf("[JBIG] will try to read %u bytes\n", clmbs_size_wh - 4);
#endif
  struct jbg_dec_state sd;
  jbg_dec_init(&sd);
  size_t read;
  jbg_dec_in(&sd, clmbs + 4, clmbs_size_wh - 4, &read);
  if(read != (clmbs_size_wh - 4))
    printf(":-( error while JBIG-decoding: the image could be tampered\n");
#ifdef DEBUG
  printf("[JBIG] read %zu bytes\n", read);
  printf("[JBIG] (decoded) length: %lu / %u\n", jbg_dec_getsize(&sd), (height * width / 2) / 8);
  printf("[JBIG] (decoded) size: %lu x %lu\n", jbg_dec_getheight(&sd), jbg_dec_getwidth(&sd));
#endif
  unsigned char *lmbs_ = jbg_dec_getimage(&sd, 0);
  unsigned int lm_length;
  unsigned char *lmbs = to_bitstream_flat(lmbs_, jbg_dec_getsize(&sd), &lm_length);

#ifdef DEBUG
  checksum = 0;
  /* for(i=0; i<lm_size; ++i) */
  /*   checksum += lmbs[i]; */
  /* printf("[Huffmann] (decoded) location-map checksum: %d\n", checksum); */
  checksum = 0;
  for(i=0; i<lm_length; ++i)
    checksum += lmbs[i];
  printf("[JBIG] (decoded) location-map checksum: %d\n", checksum);
#endif

  //destroy intermediate location maps
  free(clmbs);
  free(lmbs_);


  /*
    4.2. restore the original difference values
  */
  if(verbose) printf("> restore the original difference values\n");

#ifdef DEBUG
  printf("starting to read the lm at bit: %d\n", (int)((bitstream - bitstream_)));
#endif
  bitstream += clmbs_length_wh;

#ifdef DEBUG
  printf("starting to read LSBs at bit: %d\n", (int)((bitstream - bitstream_)));
#endif
  bitstream = restore_differences(differences, height, width/2, bitstream, membership, lmbs);

#ifdef DEBUG
  checksum = 0;
  for(i=0; i<height; ++i)
    for(j=0; j<width/2; ++j)
      checksum += differences[i][j];
  printf("restored differences checksum: %d\n", checksum);
#endif


  /*
    4.3. get the payload
  */
  if(verbose) printf("> read the embedded watermark\n");
#ifdef DEBUG
  printf("starting to read the payload at bit: %d\n", (int)((bitstream - bitstream_)));
#endif

  unsigned int payload_size = to_int32(bitstream, 0);
#ifdef DEBUG
  printf("payload size: %d bits\n", payload_size);
#endif
  bitstream += 32;

  unsigned int payload_words = payload_size;
  unsigned char *payload = (unsigned char*)malloc(payload_words * sizeof(unsigned char));
  unsigned char w;
  for(i=0; i<payload_words; ++i) {
    w = 0;
    for(j=7; j>=0; --j)
      w = w | (*(bitstream++) << j);
    payload[i] = w;
  }

#ifdef DEBUG
  checksum = 0;
  for(i=0; i<payload_words; ++i)
    checksum += payload[i];
  printf("+ watermark checksum: %u\n", checksum);
  for(i=0; i<payload_words; ++i)
    printf("%c",payload[i]);
#endif
  
  


  /*
    5.1. restore the origianl image (inverse hwt)
  */
  if(verbose) printf("> restore the original image (inverse transform)\n");
  ihwt_bitmap(restored_data, differences, average, height, width);

#ifdef DEBUG
  checksum = 0;
  for(i=0; i<height; ++i)
    for(j=0; j<width; ++j)
      checksum += restored_data[i][j];
  printf("restored bitmap checksum: %d\n", checksum);
#endif


  /*
    5.2. authenticate
  */
  if(verbose) printf("> authentication\n");
  //payload message (md5 checksum)
  md5_state_t md5_state;
  unsigned char digest[16];
  md5_init(&md5_state);
  for(i=0; i<height; ++i)
    md5_append(&md5_state, restored_data[i], width);
  md5_finish(&md5_state, digest);

  //check
  int flag = 0;
  for(i=0; i<16; ++i)
    if(digest[i] != payload[i]) {
      flag = 1;
      break;
    }
  if(!flag)
    printf(":-) watermark match the payload: the image has not been tampered\n");
  else
    printf(":-( watermark doesn't match the payload: the image has been tampered\n");


  //clean-up
  for(i=0; i<height; ++i) {
    free(differences[i]);
    free(membership[i]);
    free(average[i]);
  }
  free(differences);
  free(membership);
  free(average);
  free(lmbs);
  free(payload);
  free(bitstream_);

  return 0;
}
